import os
import numpy as np
import pandas as pd
import torch
import matplotlib.pyplot as plt
from scipy.ndimage import gaussian_filter1d  # 用于平滑处理
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score  # 用于计算指标
from test8 import WearModel, preprocess_channel, load_data  # 确保 test11.py 在同目录或 PYTHONPATH 中

def exponential_smoothing(data, alpha=0.3):
    """
    使用指数平滑法对数据进行平滑处理
    :param data: 输入数据列表
    :param alpha: 平滑系数，范围为 [0, 1]，值越大越依赖当前值，越小越平滑
    :return: 平滑后的数据列表
    """
    smoothed = [data[0]]  # 初始化第一个值
    for i in range(1, len(data)):
        smoothed_value = alpha * data[i] + (1 - alpha) * smoothed[-1]
        smoothed.append(smoothed_value)
    return smoothed

def predict_tool_wear(model, data, device):
    """
    使用模型预测刀具磨损
    :param model: 训练好的模型
    :param data: 预处理后的数据
    :param device: 设备（CPU 或 GPU）
    :return: 预测结果列表
    """
    model.eval()
    predictions = []
    with torch.no_grad():
        for sample in data:
            sample_tensor = torch.FloatTensor(sample).unsqueeze(0).to(device)  # (1, 7, window_size)
            output = model(sample_tensor)
            predictions.append(output.item())
    return predictions

def main():
    # 配置
    data_root = './Code/data'
    # model_path = './logs/best_model1_20250525_103955.pth'  # 替换为实际模型路径11
    # model_path = './logs/best_model1_20250523_230800.pth'  # 替换为实际模型路径10
    # model_path = './logs/best_model1_20250524_214541.pth'  # 替换为实际模型路径10
    # model_path = './logs/best_model_20250525_235316.pth'  # 替换为实际模型路径6
    # model_path = './logs/best_model_20250526_003640.pth'  # 替换为实际模型路径4
    # model_path = './logs/best_model_20250526_010354.pth'  # 替换为实际模型路径5
    # model_path = './logs/best_model1_20250526_022725.pth'  # 替换为实际模型路径11
    # model_path = './logs/best_model1_20250526_093949.pth'  # 替换为实际模型路径13
    # model_path = './logs/best_model1_20250526_101339.pth'  # 替换为实际模型路径14
    # model_path = './logs/best_model0_20250526_111622.pth'  # 替换为实际模型路径8
    # model_path = './logs/best_model1_20250715_214345.pth'  # 替换为实际模型路径11
    model_path = './logs/best_model_20250717_151919.pth'  # 替换为实际模型路径15
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 加载模型
    model = WearModel()
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.to(device)
    print(f"Loaded model from {model_path} on {device}")

    # 加载数据和标签
    print("Loading data and labels...")
    data, true_values = load_data(data_root)  # 使用 test11.py 中的 load_data 函数
    print(f"Loaded {len(data)} samples with labels.")

    # 预测
    print("Predicting tool wear...")
    predictions = predict_tool_wear(model, data, device)


    # 平滑处理
    # smoothed_predictions = gaussian_filter1d(predictions, sigma=2)  # 使用高斯滤波进行平滑
    #  # 平滑处理（使用指数平滑法）
    smoothed_predictions = exponential_smoothing(predictions, alpha=0.5)  # 调整 alpha 控制平滑程度

    # 计算指标
    mae = mean_absolute_error(true_values, predictions)
    rmse = np.sqrt(mean_squared_error(true_values, predictions))
    r2 = r2_score(true_values, predictions)

    print(f"MAE: {mae:.4f}")
    print(f"RMSE: {rmse:.4f}")
    print(f"R²: {r2:.4f}")

    # 绘制预测结果
    plt.figure(figsize=(10, 6))
    plt.plot(range(1, len(predictions) + 1), predictions, marker='o', color='blue', label=f'Raw Prediction')
    plt.plot(range(1, len(predictions) + 1), smoothed_predictions, color='red', label='Smoothed Prediction', linewidth=2)
    plt.plot(range(1, len(true_values) + 1), true_values, color='green', linestyle='--', label='Actual Labels')
    plt.title(f"Prediction of Tool 6 Wear")
    plt.ylabel("RUL")
    plt.xlabel("Used tool life")
    plt.legend()
    plt.grid(False)

    # 去除上和右边框
    ax = plt.gca()
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)

    plt.savefig(f'tool_prediction_with_actual_labels.png')
    plt.show()

if __name__ == '__main__':
    main()